#pragma once

#include "base_types.h"
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>

// Socket utility functions
inline int SetNonBlocking(int sockfd) {
    int flags = fcntl(sockfd, F_GETFL, 0);
    if (flags < 0) return flags;
    return fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);
}

inline int SetReuseAddr(int sockfd) {
    int optval = 1;
    return setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
}

inline int SetReusePort(int sockfd) {
    int optval = 1;
    return setsockopt(sockfd, SOL_SOCKET, SO_REUSEPORT, &optval, sizeof(optval));
}

inline int SetKeepAlive(int sockfd, bool enable = true) {
    int optval = enable ? 1 : 0;
    return setsockopt(sockfd, SOL_SOCKET, SO_KEEPALIVE, &optval, sizeof(optval));
}

inline int CreateTcpSocket() {
    int sockfd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
    if (sockfd < 0) return ERR_SOCKET;
    
    if (SetNonBlocking(sockfd) < 0) {
        close(sockfd);
        return ERR_SOCKET;
    }
    
    if (SetReuseAddr(sockfd) < 0) {
        close(sockfd);
        return ERR_SOCKET;
    }
    
    return sockfd;
}

inline int BindSocket(int sockfd, const char* ip, uint16_t port) {
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    
    if (inet_pton(AF_INET, ip, &addr.sin_addr) <= 0) {
        return ERR_BIND;
    }
    
    if (bind(sockfd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        return ERR_BIND;
    }
    
    return ERR_SUCCESS;
}

inline int ListenSocket(int sockfd, int backlog = 128) {
    if (listen(sockfd, backlog) < 0) {
        return ERR_LISTEN;
    }
    return ERR_SUCCESS;
}

inline int AcceptSocket(int sockfd, struct sockaddr_in* peer) {
    socklen_t peer_len = sizeof(struct sockaddr_in);
    int newfd = accept(sockfd, (struct sockaddr*)peer, &peer_len);
    
    if (newfd < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            return ERR_SUCCESS; // No more connections
        }
        return ERR_ACCEPT;
    }
    
    if (SetNonBlocking(newfd) < 0) {
        close(newfd);
        return ERR_ACCEPT;
    }
    
    return newfd;
}

inline int ConnectSocket(int sockfd, const char* ip, uint16_t port) {
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    
    if (inet_pton(AF_INET, ip, &addr.sin_addr) <= 0) {
        return ERR_SOCKET;
    }
    
    int ret = connect(sockfd, (struct sockaddr*)&addr, sizeof(addr));
    if (ret < 0) {
        if (errno == EINPROGRESS) {
            return ERR_SUCCESS; // Non-blocking connect in progress
        }
        return ERR_SOCKET;
    }
    
    return ERR_SUCCESS;
}

inline const char* GetPeerIp(int sockfd, char* buffer, size_t size) {
    struct sockaddr_in addr;
    socklen_t len = sizeof(addr);
    if (getpeername(sockfd, (struct sockaddr*)&addr, &len) < 0) {
        return "unknown";
    }
    return inet_ntop(AF_INET, &addr.sin_addr, buffer, size);
}

inline uint16_t GetPeerPort(int sockfd) {
    struct sockaddr_in addr;
    socklen_t len = sizeof(addr);
    if (getpeername(sockfd, (struct sockaddr*)&addr, &len) < 0) {
        return 0;
    }
    return ntohs(addr.sin_port);
}